import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg

# 弹簧-阻尼质量块位置的控制例子

# Parameters
dt = 0.1  # Time step
t_final = 10  # Total time
N = int(t_final / dt)  # Number of time steps

g = 10

# System matrices 离散型
A = np.array([[0.9951, 0.09738], [-0.09738, 0.9464]])
B = np.array([[0.0049], [0.09738]])

# Weight matrices for LQR
Q = np.array([[100, 0], [0, 1]])
R = np.array([[0.1]])

# Solve Discrete-time Algebraic Riccati Equation (DARE) for P
P = scipy.linalg.solve_discrete_are(A, B, Q, R)

# Compute LQR gain matrix K
K = np.linalg.inv(B.T @ P @ B + R) @ (B.T @ P @ A)
print(K)

# Initialize state vector
x = np.array([[0], [0]])  # Initial state [position; velocity]

# Desired trajectory
x_ref = np.zeros((2, N))
x_ref[0, :] = 1  # Reference position is a linear path from 0 to 10

# Store trajectory
x_history = list()
u_history = list()

# Simulation loop
for i in range(N):
    # Control input using LQR
    u = -K @ (x - x_ref[:, i].reshape(2, 1))
    
    # Apply control input to the system
    x = A @ x + B @ u
    
    # Record data
    x_history.append(x[0])
    u_history.append(u[0])

# Plotting results
plt.figure(figsize=(10, 5))

# Plot position
plt.subplot(2, 1, 1)
plt.plot(x_history, label='Car Position')
plt.plot(x_ref[0, :], label='Reference Position', linestyle='dashed')
plt.xlabel('Time [s]')
plt.ylabel('Position [m]')
plt.legend()
plt.title('Position Tracking')

# Plot control input
plt.subplot(2, 1, 2)
plt.plot(u_history, label='Control Input (Acceleration)')
plt.xlabel('Time [s]')
plt.ylabel('Control Input [m/s^2]')
plt.legend()
plt.title('Control Input')

plt.tight_layout()
plt.show()
